(*
* Management interface
*)

open Printf
open ExtList

open Bmap
open Constants
open Types
open Randif
open ExtList
open Array 
open Random ;;

type bsside =
  | Asset
  | Liability
  | Equity
and pnum = 
  | PCapitalRat 
  | PLiqRat 
  | PRate 
  | PDefault 
  | PPrincipal 
  | PMaturity 
  | PBankCash 
  | PCashRate 
  | PSum of pnum * pnum
  | PDiff of pnum * pnum
  | PProd of pnum * pnum
  | PDiv of pnum * pnum
  | PBSGap of pnum * pnum (* assets - liabilities at pnum maturity *) 
  | PBSSum of bsside * policy
  | PNum of float
  | PNullNum 
and pdirection =
  | PGT
  | PLT
and policy = 
  | PNullPolicy 
  | PAccept 
  | PReject
  | PNot of policy
  | PIf of policy * policy * policy
  | PCall of int
  | PNumTest of pnum * pdirection * pnum


let print_pol p =
  match p with 
  | PNullPolicy -> "Z"
  | PAccept -> "*"
  | PReject -> "."
  | (PCall n) -> sprintf "%d" n
  | _ -> "X"

let mutate_num pd pn =
  match pn with
  | (PNum n) ->
      let rn = getDNum pd in
      (PNum (n +. rn))
  | _ -> raise NOT_IMPLEMENTED
  

let mutate_policy pd pn =
  match pn with
  | PNumTest (item,test,ratio) ->
      let newratio = mutate_num pd ratio in
      PNumTest (item,test,newratio)
  | _ -> raise NOT_IMPLEMENTED

(*
* Table management
*)

type s_management = {
    lenderPolicy : policy ;
    borrowerPolicy : policy ;
    dividendPolicy : pnum ;
    registers : policy list }

type t_management = {
    lender_table : policy array;
    borrower_table : policy array;
    rate_breaks : float list;
    rate_extent : int;
    maturity_breaks : float list;
    maturity_extent : int;
    lender_tests : policy list;
    borrower_tests : policy list;
    table_div_rate : pnum}

type v_management = {
    lender_vec : float array;
    borrower_vec : float array;
    vec_maturity_breaks : float array;
    vec_maturity_extent : int;
    vec_div_rate : pnum}

type b_management = {
    accepted_lenders : bmap ;
    rejected_lenders : bmap;
    accepted_borrowers : bmap ;
    rejected_borrowers : bmap ;
    bayes_dividend_rate : pnum  }

type management =
  | SimpleM of s_management
  | TableM of t_management
  | VectorM of v_management
  | BayesM of b_management
  | NullM

type mutator = {
    mutNum : probabilityDistribution ;
    mutTree : probabilityDistribution ;
    mutTableCross : probabilityDistribution ;
    mutTablePoint : probabilityDistribution }

let lkupRegister {registers = r} n = List.nth r n

let lookupDir d ( x : float)  (y: float) = 
  match d with
  | PGT -> x > y
  | PLT -> x <= y

let empty_bayes_management = 
  BayesM { accepted_lenders = Bmap.empty cs_max prof_thresh;
	   rejected_lenders = Bmap.empty cs_max prof_thresh;
	   accepted_borrowers = Bmap.empty cs_max prof_thresh;
	   rejected_borrowers = Bmap.empty cs_max prof_thresh;
	   bayes_dividend_rate = PNum 0.02 }

let cp0 
    { lenderPolicy = lp1 ; borrowerPolicy = dp1 ;
      dividendPolicy = divp1 ; registers = reg1 }
    { lenderPolicy = lp2 ; borrowerPolicy = dp2 ;
      dividendPolicy = divp2 ; registers = reg2 } =
  let lp' = if Random.bool () then lp1 else lp2 in
  let dp' = if Random.bool () then dp1 else dp2 in 
  let divp' = if Random.bool () then divp1 else divp2 in
  let reg' = reg1 
  in { lenderPolicy = lp' ; borrowerPolicy = dp' ;
       dividendPolicy = divp' ; registers = reg' }

let new_raw_table rb mb =
  Array.make (((List.length rb) + 1) * ((List.length mb)+ 1)) PNullPolicy

let new_table_management ratebreaks matbreaks = 
  { rate_breaks = ratebreaks;
    maturity_breaks = matbreaks;
    rate_extent = (List.length ratebreaks) + 1; (* one greater than the maximum index *)
    maturity_extent = (List.length matbreaks) + 1; (* one greater than the maximum index *)
    lender_table = new_raw_table ratebreaks matbreaks;
    borrower_table = new_raw_table ratebreaks matbreaks;
    lender_tests = [];
    borrower_tests = [];
    table_div_rate = (PNum 0.05) }

(*
let lender_table_extent mgt = (mgt.rate_breaks_length)
let maturity_table_extent mgt = (mgt.maturity_breaks_length)
*)

let table_index mgt rate_interval_ix maturity_interval_ix =
  let ret = (rate_interval_ix * (mgt.maturity_extent) + maturity_interval_ix ) in
  (* printf "Looking up %d\n" ret ; *)
  ret

let set_table_management party tm pol rix mix =
  let ix = (table_index tm rix mix) in
  (* printf "%d\n" ix ; *)
  match party with
  | Lender -> Array.set tm.lender_table ix pol
  | Borrower -> Array.set tm.borrower_table ix pol

let get_table_management party tm rix mix =
  let ix = (table_index tm rix mix) in
  (* printf "%d\n" ix ; *)
  match party with
  | Lender -> Array.get tm.lender_table ix 
  | Borrower -> Array.get tm.borrower_table ix 

let proximate_entries party tm rix mix =
  let get_pentry party tm r m =
    if (r < 0) || (m < 0) || (r >= tm.rate_extent) || (m >= tm.maturity_extent)
    then None
    else Some (get_table_management party tm r m)
  in
  [ get_pentry party tm (rix+1) mix ;
    get_pentry party tm (rix-1) mix ;
    get_pentry party tm (rix) (mix + 1) ;
    get_pentry party tm (rix) (mix - 1)  ]

let proximate_entries_test t t0 =
  match t with
  | _ when List.mem None t -> false
  | _ when not (List.for_all (fun x -> (x = (Some t0))) t) -> true 
  | [ Some PAccept;Some PAccept;Some PAccept;Some PAccept ] -> false
  | [ Some PReject;  Some PReject;  Some PReject;  Some PReject ] -> false
  | _ -> true

module Management =
struct
  type t = management
  let compare = compare
end ;;

let list_crossover l1 l2 =
  let rec me1 a1 a2 =
    match (a1,a2) with
    | ([],_) -> []
    | (_,[]) -> []
    | ((m1 :: m1s),(m2 :: m2s)) -> m1 :: (me2 m1s m2s)
  and me2 a1 a2 =
    match (a1,a2) with
    | ([],_) -> []
    | (_,[]) -> []
    | ((m1 :: m1s),(m2 :: m2s)) -> m1 :: (me1 m1s m2s)
  in
  me1 l1 l2


let summarize_bmaps blist =
  let assbm func ls =
    let bmaps = List.map func ls in
    List.fold_left Bmap.merge (List.hd bmaps) (List.tl bmaps)
  in
  {
    accepted_lenders = assbm (fun x -> x.accepted_lenders) blist;
    rejected_lenders = assbm (fun x -> x.rejected_lenders) blist;
    accepted_borrowers = assbm (fun x -> x.accepted_borrowers) blist;
    rejected_borrowers = assbm (fun x -> x.rejected_borrowers) blist;
    bayes_dividend_rate = PNum 0.1  }

let merge_bayes_management (BayesM b1) (BayesM b2) =
  BayesM { accepted_lenders = Bmap.merge b1.accepted_lenders b2.accepted_lenders;
	   rejected_lenders = Bmap.merge b1.rejected_lenders b2.rejected_lenders;
	   accepted_borrowers = Bmap.merge b1.accepted_borrowers b2.accepted_borrowers;
	   rejected_borrowers = Bmap.merge b1.rejected_borrowers b2.rejected_borrowers;
	   bayes_dividend_rate = PNum 0.1 }
 
let print_bayes_management (BayesM b) =
   (Bmap.bmap_print 1 b.accepted_lenders)^
   (Bmap.bmap_print 2 b.rejected_lenders)^
   (Bmap.bmap_print 3 b.accepted_borrowers)^ 
   (Bmap.bmap_print 4 b.rejected_borrowers)


let table_crossover t1 t2 =
  match (t1,t2) with
  | ((TableM t1),(TableM t2))  ->
  (let rval nt tab1 tab2 ix =
    let vix = Random.int 2 in
    if vix == 0
    then Array.set nt ix (Array.get tab1 ix)
    else Array.set nt ix (Array.get tab2 ix)
  in 
  let tabx ( v1 : policy array) ( v2 : policy array ) = 
    let nt = new_raw_table t1.rate_breaks t1.maturity_breaks in
    List.iter (rval nt v1 v2) (0 -- ((Array.length v1) - 1));
    nt
  in
  let bt = tabx t1.borrower_table t2.borrower_table in
  let lt = tabx t1.lender_table t2.lender_table in
  TableM { t1 with lender_table = lt;
	   borrower_table = bt;
	   lender_tests = list_crossover t1.lender_tests t2.lender_tests;
	   borrower_tests = list_crossover t1.borrower_tests t2.borrower_tests })
  | ((VectorM v1),(VectorM v2)) ->
      (VectorM { lender_vec = Array.copy v1.lender_vec ;
		 borrower_vec = Array.copy v2.borrower_vec ;
		 vec_maturity_breaks = v1.vec_maturity_breaks ;
		 vec_maturity_extent = v1.vec_maturity_extent ;
		 vec_div_rate = v1.vec_div_rate })
  | ((BayesM b1),(BayesM b2)) ->
      (BayesM {(summarize_bmaps [b1; b2]) with 
	      bayes_dividend_rate = b1.bayes_dividend_rate })
  | (_,_) -> raise INTERNAL_ERROR


let divrate_mutate t = 
  let pd = (UniformPD (-0.015,0.015)) in
  match t with 
  | (TableM tab) -> 
  TableM { tab with 
	   table_div_rate = mutate_num pd tab.table_div_rate } 
  | (VectorM vec) ->
      VectorM { vec with
		vec_div_rate = mutate_num pd vec.vec_div_rate } 
  | (BayesM bmgt) ->
      BayesM { bmgt with
		bayes_dividend_rate = mutate_num pd bmgt.bayes_dividend_rate } 
  | (_) -> raise INTERNAL_ERROR

let table_policy_mutate t =
  match t with
  | (TableM tab) ->
  (let pd = (UniformPD (-0.015,0.015)) in
  TableM { tab with
	   lender_tests = List.map (fun t -> mutate_policy pd t) tab.lender_tests;
	   borrower_tests = List.map (fun t -> mutate_policy pd t) tab.borrower_tests })
  | _ -> raise INTERNAL_ERROR


let table_indices tab =
  let ri = (0 -- (tab.rate_extent - 1)) in
  let mi = (0 -- (tab.maturity_extent - 1)) in
  let rec z2 i m =
    match m with 
    | [] -> []
    | m1 :: ms -> (i, m1) :: (z2 i ms)
  in
  List.concat (List.map (fun r -> z2 r mi) ri)

type vmutator = {
    vmutmin : float;
    vmutmax : float;
    vmutstep : float }


let adj_val vmut v =
  let adjv = if Random.bool() 
             then v -. vmut.vmutstep
             else v +. vmut.vmutstep in
  if adjv < vmut.vmutmin 
  then vmut.vmutmin
  else (if adjv > vmut.vmutmax
  then vmut.vmutmax
  else adjv )
      

(* Makes a copy of vecarr *)
let vec_mutate vmut vecarr =
  let v2 = Array.copy vecarr in
  let len = Array.length vecarr in
  let rix = Random.int len in
  Array.set v2 rix (adj_val vmut vecarr.(rix));
  v2
  

let point_mutate t = 
  match t with 
  | (TableM tab) -> 
  (let mval party t0 t rix mix =
    let ix = table_index tab rix mix in
    let this = get_table_management party tab rix mix in 
    let rc = Random.int 100 in
    if rc < table_point_mutate_bound &&
      proximate_entries_test (proximate_entries party tab rix mix) this
    then
      let vix = Random.int 2 in
      let r = if vix == 0
              then PAccept
              else PReject
      in
      Array.set t ix r
    else Array.set t ix (Array.get t0 ix)
  in
  let tabm p mt = 
    let nt = new_raw_table tab.rate_breaks tab.maturity_breaks in
    List.iter (fun (r,m) -> mval p mt nt r m) (table_indices tab);
    nt
  in
  let r = TableM { tab with lender_table = tabm Lender tab.lender_table ;
		   borrower_table = tabm Borrower tab.borrower_table } in
  table_policy_mutate ( divrate_mutate r))
  | (VectorM vectm) -> 
      let vmut = { vmutmin = 0.0;
		       vmutmax = 0.1;
		       vmutstep = 0.05} in
      let r = { vectm with
		lender_vec = vec_mutate vmut vectm.lender_vec; 
		borrower_vec = vec_mutate vmut vectm.borrower_vec } in
      (divrate_mutate (VectorM r))
  | (BayesM bm) -> (divrate_mutate (BayesM bm))
  | _ -> raise INTERNAL_ERROR


let mutate_management ms = 
  let nm = List.map point_mutate ms in
  nm

(*
let mutate_management ms =
  let nm = List.map point_mutate ms in
  let (n1,n2) = List.split_nth ((List.length nm) / 2) nm in
  let hl = List.map2 table_crossover n1 n2 in
  hl @ n1
*)

let divrate b t =
  match t with
  | (TableM { table_div_rate = dr}) ->
      (match dr with
      | (PNum n) -> n
      | _ -> raise BAD_MANAGEMENT)
  | (VectorM { vec_div_rate = dr}) ->
      (match dr with
      | (PNum n) -> n
      | _ -> raise BAD_MANAGEMENT)
  | (BayesM { bayes_dividend_rate = dr}) ->
      (match dr with
      | (PNum n) -> n
      | _ -> raise BAD_MANAGEMENT)
  | _ -> raise INTERNAL_ERROR


let print_table m party =
  let pmats ms =
    List.fold_left (fun s n -> s ^ (sprintf " %4.2f " n)) "" ms
  in
  let matixs = (0 -- (List.length m.maturity_breaks))  in
  let prow st rix =
    (sprintf ">%4.2f " st) ^ 
    (List.fold_left (fun s mix -> s ^ (sprintf " %4s " (print_pol (get_table_management party m rix mix)) )) "" matixs)
  in
  let rats = List.rev (List.combine (0. :: m.rate_breaks) (0 -- (List.length m.rate_breaks))) in
  "           " ^ (pmats m.maturity_breaks) ^ "\n"
  ^ (List.fold_left (fun s (st,ix) -> s ^ (prow st ix) ^ "\n") "" rats  )

let print_management t =
  match t with
  | (TableM 
			      ({ lender_table = lt;
				borrower_table = bt;
				rate_breaks = rb;
				maturity_breaks = mb;
				table_div_rate = dr;
			        lender_tests = lp;
			        borrower_tests = bp } as tm)) ->
  ("Lender (wholesale funding):\n" ^ (print_table tm Lender) ^
  "Borrower (writing loans):\n" ^ (print_table tm Borrower) ^
  "Dividend rate: " ^ (Std.dump dr)  ^ "\n"
  ^ "Lender Tests: " ^ (Std.dump lp) ^ "\n" ^
  "Borrower Tests: " ^ (Std.dump bp) ^"\n")
  | (BayesM  b) ->
      sprintf "BayesMgt Summary #acc:rej l %6d:%6d ---  #acc:rej b %6d:%6d div %5.3f \n"
	(Bmap.length b.accepted_lenders) (Bmap.length b.rejected_lenders)
	(Bmap.length b.accepted_borrowers) (Bmap.length b.rejected_borrowers)
	(divrate 0.0 (BayesM b))
  | _ -> raise INTERNAL_ERROR

